import cv2
import numpy as np
from deep_sort_realtime.deepsort_tracker import DeepSort

def track_objects(frame, detections, colors, tracker):
    """
    使用 DeepSORT 进行多目标跟踪并绘制结果
    :param frame: 当前帧图像
    :param detections: 检测结果列表，每个元素为 (bounding_box, confidence, class_label)
    :param colors: 颜色列表
    :param tracker: DeepSORT 跟踪器
    :return: 处理后的帧图像
    """
    # 获取图片尺寸
    height, width = frame.shape[:2]
    
    # 使用原始图像作为画布
    canvas = frame.copy()
    
    tracks = tracker.update_tracks(detections, frame=frame)
    
    # 情绪颜色映射 (BGR格式)
    emotion_colors = {
        'Happy': (0, 255, 127),     # 绿色
        'Sad': (255, 0, 0),         # 蓝色
        'Angry': (0, 0, 255),       # 红色
        'Surprise': (0, 255, 255),  # 黄色
        'Neutral': (169, 169, 169), # 灰色
        'Calm': (255, 255, 0),      # 青色
        'Fear': (128, 0, 128),      # 紫色
        'Disgust': (139, 69, 19),   # 棕色
        'Contempt': (0, 140, 255),  # 橙色
        'Unknown': (192, 192, 192)  # 浅灰色
    }
    
    for track in tracks:
        if not track.is_confirmed():
            continue
            
        track_id = track.track_id
        bbox = track.to_ltrb()
        x1, y1, x2, y2 = map(int, bbox)
        emotion_label = track.det_class
        
        # 转换情绪标签为中文
        emotion_chinese = {
            'Happy': '高兴',
            'Sad': '伤心',
            'Angry': '生气',
            'Surprise': '惊讶',
            'Neutral': '中性',
            'Calm': '平静',
            'Fear': '恐惧',
            'Disgust': '厌恶',
            'Contempt': '蔑视',
            'Unknown': '未知'
        }.get(emotion_label, '未知')

        # 获取对应的情绪颜色
        color = emotion_colors.get(emotion_label, (192, 192, 192))
        
        # 绘制边界框（加粗）
        cv2.rectangle(canvas, (x1, y1), (x2, y2), color, 3)
        
        # 准备标签文本
        label_text = f"{emotion_chinese}"
        
        # 计算标签大小
        font_scale = 0.8
        font_thickness = 2
        (label_width, label_height), baseline = cv2.getTextSize(
            label_text, cv2.FONT_HERSHEY_SIMPLEX, font_scale, font_thickness
        )
        
        # 确保标签位置在图像内
        label_x = max(x1, 0)
        label_y = max(y1 - 10, label_height)  # 将标签放在框的上方
        
        # 绘制标签背景（半透明）
        sub_img = canvas[label_y - label_height - baseline:label_y + baseline,
                        label_x:label_x + label_width]
        rect = np.zeros(sub_img.shape, dtype=np.uint8)
        rect[:] = color
        res = cv2.addWeighted(sub_img, 0.7, rect, 0.3, 1.0)
        canvas[label_y - label_height - baseline:label_y + baseline,
               label_x:label_x + label_width] = res
        
        # 绘制标签文本
        cv2.putText(
            canvas,
            label_text,
            (label_x, label_y),
            cv2.FONT_HERSHEY_SIMPLEX,
            font_scale,
            (255, 255, 255),  # 白色文字
            font_thickness
        )

    return canvas